// Copyright 2024 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

syntax = "proto2";

package eeg_modelling.protos;


// This message contains request parameters for waveform data and prediction
// data.  The request will be created by the frontend and then sent to the
// backend which will return the response specified by the request.
// Next ID: 15
message DataRequest {
  // Path to SSTable containing TF Examples.
  optional string tf_ex_sstable_path = 1;
  // Key for the TF Example in the SSTable.
  optional string sstable_key = 2;
  // Path to SSTable containing PredictionOutputs.
  optional string prediction_sstable_path = 3;
  // Duration of the requested chunk in seconds.
  optional int64 chunk_duration_secs = 4;
  // Relative start time of the requested chunk in seconds.
  optional int64 chunk_start = 5;
  // List of the requested channel indices.
  repeated ChannelDataId channel_data_ids = 12;
  // Value to use as low cut for high pass filter.
  optional float low_cut = 7;
  // Value to use as high cut for low pass filter.
  optional float high_cut = 8;
  // Value to use for the notch filter.
  optional float notch = 9;
  // Path to Edf.
  optional string edf_path = 10;
  // Segment index in the Edf.
  optional int64 edf_segment_index = 11;
  // List of the requested feature indices.
  repeated Feature features = 6 [deprecated = true];
  // Path to a local file containing a TF example protobuf
  optional string tf_ex_file_path = 13;
  // Path to a local file containing a PredictionsOutput protobuf
  optional string prediction_file_path = 14;
}

// This message contains the data request response.  This response will be sent
// from the backend to the frontend on request containing the requested waveform
// and prediction data.
message DataResponse {
  // Waveform chunk-wise data.
  optional WaveformChunk waveform_chunk = 1;
  // Waveform metadata.
  optional WaveformMetadata waveform_metadata = 2;
  // Prediction chunk-wise data.
  optional PredictionChunk prediction_chunk = 3;
  // Waveform chunk-wise data.
  optional PredictionMetadata prediction_metadata = 4;
}

// This message contains waveforms data specific to a time chunk.
message WaveformChunk {
  // Waveform data packed into a DataTable JSON string.
  optional string waveform_datatable = 1;
  // Indices of the channels in the data.
  repeated ChannelDataId channel_data_ids = 4;
  // Sampling frequency of data in DataTable.
  optional float sampling_freq = 3;
  // Indices of the channels in the data.
  repeated Feature features = 2 [deprecated = true];
}

// This message contains waveform metadata consistent across a file.
message WaveformMetadata {
  // Absolute start time of the file in time since epoch.
  optional double abs_start = 1;
  // Helper proto to describe annotations.
  message Label {
    // Text associated with the annotation.
    optional string label_text = 1;
    // Relative start time of the annotation.
    optional float start_time = 2;
  }
  // List of annotations.
  repeated Label labels = 2;
  // Dictionary of generated channel indices to their shorthand names.
  map<string, string> channel_dict = 3;
  // Type of waveforms data in the file.
  optional string file_type = 4;
  // Number of seconds in the entire file.
  optional float num_secs = 6;
  // Patient ID.
  optional string patient_id = 7;
  // Key used to request the SSTable values (may have been fetched if none was
  // provided).
  optional string sstable_key = 8;
  // JSON string of a DataTable with only a time axis for the entire file.
  optional string nav_timeline_datatable = 5 [deprecated = true];
}

// This message contains prediction specific to a time chunk.
message PredictionChunk {
  // Helper proto to describe attribution map.
  message AttributionMap {
    // Helper proto to describe attribution data.
    message AttributionValues {
      // List of attribution values for a channel.
      repeated float attribution = 1;
    }
    // Dictionary of channel indices to their attribution values.
    map<string, AttributionValues> attribution_map = 1;
  }
  // Attribution data for each label
  map<string, AttributionMap> attribution_data = 1;
  // The relative start time to the input file start.
  optional int64 chunk_start = 3;
  // The length of the chunk in seconds.
  optional int64 chunk_duration = 4;
  // JSON string of a DataTable with only a time axis for the chunk.
  optional string attribution_datatable = 2 [deprecated = true];
}

// This message contains prediction metadata consistent across a file.
message PredictionMetadata {
  // Helper proto to describe metadata and scores for a chunk.
  message ChunkScoreData {
    // Duration of the chunk being scored in seconds.
    optional float duration = 1;
    // Relative start time of the chunk.
    optional float start_time = 2;
    // Helper proto to describe chunk scores.
    message ScoreData {
      // Score that determines true positive/negative.
      optional float actual_value = 1;
      // Floating point score output by the model.
      optional float predicted_value = 2;
      // Probability (between 0..1) of the prediction made
      optional float prediction_probability = 3;
    }
    // Dictionary of prediction types to score data.
    map<string, ScoreData> score_data = 3;
  }
  // List of chunk score data.
  repeated ChunkScoreData chunk_scores = 1;
}

// Helper proto for time series data request ID's.
message ChannelDataId {
  // Helper proto for time series requests that refer to a single channel.
  message SingleChannel {
    // Index of the channel based on WaveformMetadata channel_dict.
    optional int64 index = 1;
  }
  // Helper proto for time series requests that refer to a subtraction between
  // two channels.  The subtraction performed is (primary - referential).
  message BipolarChannel {
    // Index of the primary channel.
    optional int64 index = 1;
    // Index of the referential channel.
    optional int64 referential_index = 2;
  }
  oneof id {
    SingleChannel single_channel = 1;
    BipolarChannel bipolar_channel = 2;
  }
}

// Helper proto to describe channel features.
message Feature {
  // Feature index (or indices for standards).
  optional string feature = 1;
}
